import time

import cv2
import numpy as np
import torch, sys, os

from recog.FOTS.data_loader.datautils import normalize_iamge
from recog.FOTS.model.raw_crnn import RAWCRNN
from recog.FOTS.utils.util import str_label_converter
from recog.train_raw_crnn import LpDetSetWithColor


def get_crnn_model():
    file_dir = os.path.dirname(os.path.abspath(__file__))
    cur_dir = os.getcwd()
    os.chdir(file_dir)

    model = RAWCRNN(40, 3, 69, 256)
    # model.load_state_dict(torch.load("/project/train/models/CRNN91_l0.00064566_tc1.000_vc0.938_tcc1.000_vcc1.000"))
    model.load_state_dict(torch.load("/project/train/models/CRNN18_l0.00767774_tc0.989_vc0.993_tcc1.000_vcc1.000"))
    model.eval()
    if torch.cuda.is_available():
        model = model.cuda()

    os.chdir(cur_dir)

    return model


def get_lp_str(bgr_im, model):
    if bgr_im.shape[0] != 40 or bgr_im.shape[1] != 120:
        bgr_im = cv2.resize(bgr_im, (120, 40))
    im_tensor = normalize_iamge(bgr_im, aug=False)
    im_tensor = torch.from_numpy(im_tensor).cuda().float()
    im_tensor = im_tensor.permute(2, 0, 1)
    im_tensor = im_tensor.unsqueeze(0)

    t = time.time()
    out, logits = model(im_tensor)  # T B C
    lp_length = torch.tensor([8])
    print("cost lp recog:{}".format(time.time() - t))

    with torch.no_grad():
        pred = [out, torch.fill_(torch.ones_like(lp_length), out.shape[0]).int()]

        pred_tensor = torch.softmax(pred[0], dim=-1).argmax(dim=-1).t().detach().cpu()
        decode_pred = str_label_converter.decode(pred_tensor, torch.tensor(pred[1].int()))
        if isinstance(decode_pred, list) is False:
            decode_pred = [decode_pred]
        decode_pred = str_label_converter.remove_duplicate(decode_pred)
    return decode_pred[0], LpDetSetWithColor.clsID2color(int(logits.argmax(dim=1)[0]))


if __name__ == '__main__':

    model = get_crnn_model()

    root = "/home/leo/Downloads/data_set/lp_reg/val"
    for file in os.listdir(root):
        if ".png" in file:
            bgr_im = cv2.imread(os.path.join(root, file))
            lp_str, color = get_lp_str(bgr_im, model)

            print(lp_str)
            cv2.imshow("x", bgr_im)
            cv2.waitKey(0)
